#!/usr/bin/env python

import logging
import math
import os
import os.path
import paramiko
import signal
import subprocess
import sys
import time
import multiprocessing

#Local imports
import virtmanager

class VirtLXC:
	"""VM-script to work with both local and remote LXC"""
	def __init__(self, total_nodes, libvirt_xml, remote_folder, remote_ips, experiment_duration, initial_wait, shutdown_wait, remote_act_folders, mtu, log='INFO', memory='500000', taskset='0'):
		if log.upper() == 'INFO':
			lev = logging.INFO
		elif log.upper() == 'DEBUG':
			lev = logging.DEBUG
		else:
			lev = logging.WARN
		logging.basicConfig(level=lev)
		self.log = logging.getLogger(self.__class__.__name__)
		
		self.nodes = int(total_nodes)
		self.memory = int(memory)
		self.mtu = int(mtu)
		self.taskset = taskset
		
		self.xml = libvirt_xml
		self.vm_name = 'lxc_vm_{name}'
		#self.vm_config = 'lxc_temp_conf_{nr}.xml'
		self.vm_config = 'lxc_temp_conf_{nr}.conf'
		self.pid_file = 'ssh_mobiemu_pid_file_{nr}.pid'
		
		self.ips = map(str.strip, remote_ips.split(','))
		self.real_folders = dict(zip(self.ips, map(os.path.abspath, map(str.strip, remote_act_folders.split(',')))))
		self.remote = remote_ips.strip() != '127.0.0.1'
		if self.remote:
			self.uris = dict(zip(self.ips, ['lxc+ssh://{ip}/'.format(ip=i) for i in self.ips]))
			self.ssh_connections = dict(zip(self.ips, self.__open_connections(self.ips)))
		else:
			self.uris = {}
			self.uris[self.ips[0]] = 'lxc:///'
		self.log.debug('Uris: %s', self.uris)
		self.remote_folder = remote_folder
		self.duration = int(experiment_duration) + int(initial_wait) #+ int(shutdown_wait)
		self.node_split = self.__node_split(self.ips, self.nodes)
		self.log.debug('Node split: %s', self.node_split)
		
	def __open_connections(self, ips):
		"""Open an ssh connection to all remote hosts, this should ensure that 
		we only open one connection per host"""
		result = []
		for ip in ips:
			client = paramiko.SSHClient()
			client.load_system_host_keys()
			client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
			client.connect(ip)
			result.append(client)
		return result
		
	def __close_connections(self):
		"""Close all ssh connections if there were any opened"""
		if self.ssh_connections:
			for ip in self.ssh_connections:
				try:
					self.ssh_connections[ip].close()
				except:
					self.log.exception('An error occured while trying to turn of the ssh connection to ip "%s"', ip)
					
	def close(self):
		"""Close all open resources used by this class"""
		self.__close_connections()
	
	def copy_folder(self, src_folder):
		"""Copy folder 'src_folder' to the execution folder on either local machine
		or all remote machines"""
		if not self.remote:
			dest = self.real_folders['127.0.0.1']
			if os.path.abspath(src_folder) != os.path.abspath(dest):
				self.log.debug('Copying local folder from "%s" to "%s"', src_folder, dest) 
				subprocess.call(['cp', '-r', src_folder, dest])
		else:
			for ip in self.real_folders.keys():
				self.log.debug('Copying remote folder from "%s" to "%s"', src_folder, '{}:{}'.format(ip, self.real_folders[ip]))
				subprocess.check_call(['scp', '-r', src_folder, '{}:{}'.format(ip, self.real_folders[ip])])

	def copy_file(self, src_file):
		"""Copy file to the execution folder on either local or remote machine"""
		if not self.remote:
			dest = self.real_folders['127.0.0.1']
			if os.path.abspath(src_file) != os.path.abspath(dest + os.path.basename(src_file)):
				self.log.debug('Copying local file from "%s" to "%s"', src_file, dest)
				subprocess.call(['cp', src_file, dest])
		else:
			for ip in self.real_folders.keys():
				self.log.debug('Copying remote file from "%s" to "%s"', src_file, '{}:{}'.format(ip, self.real_folders[ip]))
				subprocess.check_call(['scp', src_file, '{}:{}'.format(ip, self.real_folders[ip])])
				
	def __start_libvirt_domains(self):
		"""Start all libvirt domains, this method relies upon the previous 
		creations of and defining of libvirt configs"""
		
		#Function to use with libvirt, but as libvirt currently doesn't support
		#OLSR/BMF because it doesn't allow us to write to /dev/net/tun
		managers = self.__create_managers()
		for man in managers:
			man.start_domain()
	
	def __create_lxc_nodes(self):
		"""Create lxc nodes using regular lxc-utils methods, uses ssh for remote"""
		if self.remote:
			for i in range(1, self.nodes + 1):
				#Copy configuration to remote machines
				self.copy_file(self.vm_config.format(nr=i))
			for ip in self.node_split.keys():
				client = self.ssh_connections[ip]
				#client.exec_command('cd {path}'.format(path=self.real_folders[ip]))
				for node in self.node_split[ip]:
					command = ['cd',self.real_folders[ip], '&&', 'lxc-create','-n',self.vm_name.format(name=node),'-f',self.vm_config.format(nr=node)]
					client.exec_command(' '.join(command))
		else:
			for i in range(1, self.nodes + 1):
				self.log.debug('Creating local lxc node%s', i)
				command = ['lxc-create','-n',self.vm_name.format(name=i),'-f',self.vm_config.format(nr=i)]
				subprocess.check_call(command)

	def create_devices(self):
		"""Create the necessary devices, expanding arp table and generates 
		configuration for lxc"""
		self.__generate_lxc_conf()
		self.__expand_arp()
		self.__create_devices()
		#We must init taps and bridges before this call:
		try:
			#self.__start_libvirt_domains()
			self.__create_lxc_nodes()
		except:
			self.log.exception('Problem starting domains')
			self.__remove_devices()
			raise
		self.log.debug('Done creating devices!')
		
	def __shutdown_libvirt(self):
		"""Shut down all libvirt domains and undefine them. This method will shutdown
		all domains without throwing an error, but all errors will be logged"""
		managers = self.__create_managers()
		for man in managers:
			try:
				man.shutdown(undefine=True)
			except:
				self.log.exception('Domain threw exception while shuting down')
				continue
				
	def __remove_lxc(self):
		"""Remove LXC nodes using lxc-utils"""
		self.log.debug('Removing LXC specific devices')
		if self.remote:
			for ip in self.node_split.keys():
				client = self.ssh_connections[ip]
				for node in self.node_split[ip]:
					commands = [['lxc-stop','-n',self.vm_name.format(name=node)], ['lxc-kill','-n',self.vm_name.format(name=node)], ['lxc-destroy','-n',self.vm_name.format(name=node)]]
					for cmd in commands:
						self.__call(client, cmd)
		else:
			for i in range(1, self.nodes + 1):
				self.log.debug('Creating local lxc node%s', i)
				commands = [['lxc-stop','-n',self.vm_name.format(name=i)], ['lxc-kill','-n',self.vm_name.format(name=i)], ['lxc-destroy','-n',self.vm_name.format(name=i)]]
				map(subprocess.call, commands)
				
	def __call(self, client_connection, command):
		"""Helper method to use when executing a command on a ssh connection
		this method will block untill the call is done and return the exit
		code"""
		stdin, stdout, stderr = client_connection.exec_command(' '.join(command))
		return stdout.channel.recv_exit_status()

	def remove_devices(self):
		"""Remove all running LXC instances and remove devices both locally and remotly"""
		try:
			#self.__shutdown_libvirt()
			self.__remove_lxc()
		except:
			self.log.exception('Exception occured, nodes will not be shutdown!')
		finally:
			self.__remove_devices()

	def __run_libvirt_commands(self, start_script, experiment, settings_file):
		"""Run start script on libvirt domains"""
		managers = self.__create_managers()
		for node, man in enumerate(managers):
			node += 1
			self.log.debug('Starting node %s', node)
			man.run_command(['ifconfig', 'eth0', '10.0.0.{ip}'.format(ip=node), 'up'])
			man.run_command(['export', 'PATH=/usr/bin:/bin'])
			man.run_command(['cd', self.remote_folder])
			man.run_command(['./' + start_script, str(node), experiment, settings_file])
			#Fugly hack:
			man.stream.finish()
			man.stream = None
			
	def __start_lxc(self, start_script, experiment, settings_file):
		"""Start lxc nodes using lxc-utils"""
		self.log.debug('Starting up LXC nodes')
		if self.remote:
			for i in range(1, self.nodes + 1):
				self.copy_file(self.vm_config.format(nr=i))
			for ip in self.node_split.keys():
				client = self.ssh_connections[ip]
				#client.exec_command('cd {path}'.format(path=self.real_folders[ip]))
				for node in self.node_split[ip]:
					client.exec_command('cd {path} && lxc-execute --name {name} ./{script} {nr} {exp} {sett} > /dev/null 2>&1 &'.format(name=self.vm_name.format(name=node), path=self.real_folders[ip], script=start_script, nr=node, exp=experiment, sett=settings_file))
		else:
			for i in range(1, self.nodes + 1):
				subprocess.Popen(['lxc-execute', '--name', self.vm_name.format(name=i), './' + start_script, str(i), experiment, settings_file], close_fds=True)
	
	def start_nodes(self, start_script, experiment, settings_file):
		"""Start lxc instances"""
		#self.__run_libvirt_commands(start_script, experiment, settings_file)
		self.__start_lxc(start_script, experiment, settings_file)
		#time.sleep(self.duration)

	def move_results(self, dest):
		"""Move results from execution folder to local 'dest' folder"""
		if not self.remote:
			src = os.path.abspath(self.real_folders['127.0.0.1'] + '/dumps')
			if src != os.path.abspath(dest):
				subprocess.check_call(['cp', '-r', src, dest])
		else:
			for ip in self.real_folders.keys():
				subprocess.check_call(['scp', '-r', '{}:{}'.format(ip, self.real_folders[ip] + '/dumps/.'), dest])
				self.__remove_remote_folder(ip, self.real_folders[ip])
	
	def __remove_remote_folder(self, remote_ip, remote_folder):
		"""Remove a content of a remote folder"""
		self.log.debug('Removing remote folder at ip: %s, removing folder "%s"', remote_ip, remote_folder)
		client = self.ssh_connections[remote_ip]
		self.__call(client, ['rm', '-r', remote_folder])
		self.__call(client, ['mkdir', remote_folder])
				
	def __node_split(self, ips, nodes):
		'''This method will split the number of nodes equally between the
		different ips'''
		op = round
		length = len(ips)
		nr = int(op(float(nodes) / length))
		#print nr
		node_split = dict(zip(ips, [[]] * length))
		i = 1
		for ip in node_split.keys()[:-1]:
			node_split[ip] = range(i, i+nr)
			i += nr
		node_split[node_split.keys()[-1]] = range(i, nodes + 1)
		return node_split
	
	def __generate_lxc_conf(self):
		"""Generate lxc configurations. This method uses a template and replaces
		certain strings to create a new configuration file"""
		self.log.debug('Generating LXC configs')
		with open(self.xml, 'r') as f:
			xmldesc = f.read()
		for ip in self.node_split:
			for node in self.node_split[ip]:
				newxml = xmldesc.replace('@name@', self.vm_name.format(name=node))
				newxml = newxml.replace('@mem@', str(self.memory))
				newxml = newxml.replace('@source@', self.real_folders[ip])
				newxml = newxml.replace('@target@', self.remote_folder)
				newxml = newxml.replace('@br@', 'br{}'.format(node))
				newxml = newxml.replace('@ip_address@', '10.0.0.{ip}'.format(ip=node))
				newxml = newxml.replace('@netmask@', '10.0.0.255')
				newxml = newxml.replace('@mac@', '00:16:3e:00:01:%0.2X' % node)
				newxml = newxml.replace('@mtu@', str(self.mtu))
				newxml = newxml.replace('@cpus@', self.taskset)
				with open(self.vm_config.format(nr=node), 'w') as f:
					f.write(newxml)
	
	def __create_managers(self):
		"""Create all libvirt managers and return them in an Array"""
		self.log.debug('Generating VM managers')
		mans = []
		for ip in self.node_split.keys():
			for node in self.node_split[ip]:
				mans.append(virtmanager.VirtManager(self.vm_config.format(nr=node), self.uris[ip]))
		return mans
		
	def __create_devices(self):
		"""Create all local and remote devices"""
		if self.remote:
			self.log.debug('Setting up devices on remote machines')
			for ip in self.node_split.keys():
				for node in self.node_split[ip]:
					#ssh -o Tunnel=ethernet -f -w 0:0 2.2.2.2 true
					subprocess.check_call(['ssh', '-o', 'Tunnel=ethernet', '-f', '-w', '{nr}:{nr}'.format(nr=node), ip, 'true'])
			cmds = ['brctl addbr br{node_nr}', 'ifconfig tap{node_nr} mtu {mtu} up'.format(mtu=self.mtu, node_nr='{node_nr}'), 'brctl addif br{node_nr} tap{node_nr}', 'ifconfig br{node_nr} up']
			self.__run_remote(cmds)
		
		for i in range(1, self.nodes + 1):
			self.log.debug('Setting up devices on local machine for node %s', i)
			brdev = 'br{}'.format(i)
			tapdev = 'tap-{}'.format(i)
			tapdev2 = 'tap{}'.format(i)
			subprocess.check_call(['brctl','addbr',brdev])
			subprocess.check_call(['tunctl','-t',tapdev])
			#subprocess.check_call(['ifconfig',tapdev,'10.0.0.{ip}'.format(ip=i),'mtu',str(self.mtu),'promisc','up'])
			subprocess.check_call(['ifconfig',tapdev,'0.0.0.0','mtu',str(self.mtu),'promisc','up'])
			subprocess.check_call(['brctl','addif',brdev,tapdev])
			if self.remote:
				subprocess.check_call(['ifconfig',tapdev2,'0.0.0.0','mtu',str(self.mtu),'promisc','up'])
				subprocess.check_call(['brctl','addif',brdev,tapdev2])
			subprocess.check_call(['ifconfig',brdev, 'up'])
				
	def __remove_devices(self):
		"""Remove all local and remote devices"""
		if self.remote:
			self.log.debug('Removing devices on remote machines')
			cmds = ['ifconfig br{node_nr} down', 'ifconfig tap{node_nr} down','brctl delif br{node_nr} tap{node_nr}', 'brctl delbr br{node_nr}']
			self.__run_remote(cmds)
			try:
				pids = subprocess.check_output(['ps', '-C', 'ssh -o Tunnel=ethernet -f -w', '-o', 'pid='])
				if pids:
					for pid in pids.split():
						os.kill(int(pid), signal.SIGTERM)
			except subprocess.CalledProcessError:
				self.log.exception('Could not retreive pids for the ssh tunnels created, try "killall ssh"!')
			except:
				self.log.exception('An unexpected error occured when trying to terminate all ssh tunnels')
		
		for i in range(1, self.nodes + 1):
			self.log.debug('Removing devices for node %s on local machine', i)
			brdev = 'br{}'.format(i)
			tapdev = 'tap-{}'.format(i)
			tapdev2 = 'tap{}'.format(i)
			subprocess.call(['ifconfig',brdev,'down'])
			subprocess.call(['brctl','delif',brdev,tapdev])
			if self.remote:
				subprocess.call(['brctl','delif',brdev,tapdev2])
			subprocess.call(['brctl','delbr',brdev])
			subprocess.call(['ifconfig',tapdev,'down'])
			subprocess.call(['ip','link', 'delete',tapdev])
		self.__remove_links()
			
	def __remove_links(self):
		if self.remote:
			for ip in self.node_split.keys():
				client = self.ssh_connections[ip]
				stdin, stdout, stderr = client.exec_command('ip link show')
				ret = stdout.channel.recv_exit_status()
				if ret == 0:
					self.__remove_links2(stdout.channel.recv(1024*1024).split('\n'), client)
				else:
					self.log.warn('Could not remove links on remote macine %s!', ip)
		else:
			links = subprocess.Popen(['ip','link','show'],stdout=subprocess.PIPE,close_fds=True)
			self.__remove_links2(links.stdout.read().split('\n'))
	
	def __remove_links2(self, links, client=None):
		to_remove = []
		for link in links:
			if link.find('veth') > 0:
				device = link.split(':')[1].strip()
				to_remove.append(['ip', 'link', 'del', device])
		if self.remote:
			for cmd in to_remove:
				self.__call(client, cmd)
		else:
			map(subprocess.call, to_remove)
		
	def __run_remote(self, cmds):
		"""Run commands on all remote VMs"""
		for ip in self.node_split.keys():
			self.log.debug('Running command on remote machine with ip "%s"', ip)
			client = self.ssh_connections[ip]
			for node in self.node_split[ip]:
				#map(client.exec_command, map(lambda x: x.format(node_nr=node), cmds))
				for cmd in map(lambda x: x.format(node_nr=node), cmds):
					self.__call(client, cmd.split(' '))

	def __expand_arp(self):
		cmds = [['echo', '1024', '>', '/proc/sys/net/ipv4/neigh/default/gc_thresh1'], ['echo', '2048', '>', '/proc/sys/net/ipv4/neigh/default/gc_thresh2'], ['echo', '4096', '>', '/proc/sys/net/ipv4/neigh/default/gc_thresh3'], ['sysctl', '-w', 'net.core.rmem_max=8388608'], ['sysctl', '-w', 'net.core.wmem_max=8388608']]
		self.log.debug('Expanding ARP table')
		if self.remote:
			for ip in self.node_split.keys():
				client = self.ssh_connections[ip]
				for cmd in cmds:
					self.__call(client, cmd)
		else:
			map(subprocess.call, cmds)
		
def parse_environment():
	"""Parse the environment and return a dict of values that we use to initialize
	the VirtLXC"""
	keys = ['total_nodes', 'libvirt_xml', 'remote_folder', 'remote_ips', 'experiment_duration', 'initial_wait', 'shutdown_wait', 'remote_act_folders', 'memory', 'mtu', 'log', 'taskset']
	return dict(zip(keys, map(os.getenv, keys)))
	
def print_help():
	print 'This is a virtual machine script used my MobiEmu and is not meant to be run by it self'
	print
	print 'This vm-script has the following configuration options:'
	print 'All lists should be comma separated strings: "value1, value2, value3"'
	print '\t', 'libvirt_xml:', 'The path to the libvirt XML configuration file to be used as a template'
	print '\t', 'remote_folder:', 'The remote folder where the experiment should be performed.\n\t\tThis should equal the <target/> tag under <filesystem> in the libvirt XML if used'
	print '\t', 'remote_ips:', 'The IP addresses of the remote machines to be used.\n\t\tIf the local machine is to be used only "127.0.0.1" is enough'
	print '\t', 'remote_act_folders:', 'The remote folders where the actual filesystem is.\n\t\tThis should equal the <source/> tag under <filesystem> in the libvirt XML if used. These folder must be in the same order as the remote IP addresses.'
	print '\t', 'memory:', 'The amount of memory to give to each LXC node, default is 500000KiB'
	print '\t', 'log:', 'The logging level [INFO,DEBUG]'
	print '\t', 'taskset:', 'A string of CPUs to run the VMs on, please see "man taskset" for more information'
	
	
def main(args):
	if args:
		vl = VirtLXC(**parse_environment())
		
		if args[0] == 'copy_folder':
			src_folder = args[1]
			vl.copy_folder(src_folder)
		elif args[0] == 'copy_file':
			src_file = args[1]
			vl.copy_file(src_file)
		elif args[0] == 'create_devices':
			vl.create_devices()
		elif args[0] == 'remove_devices':
			vl.remove_devices()
		elif args[0] == 'start_nodes':
			script = args[1]
			exp = args[2]
			settings_file = args[3]
			vl.start_nodes(script, exp, settings_file)
		elif args[0] == 'move_results':
			dest = args[1]
			vl.move_results(dest)
		else:
			print 'No valid option selected!'
		
		vl.close()
	else:
		print_help()

if __name__ == '__main__':
	main(sys.argv[1:])